from sklearn.datasets import load_iris
from sklearn.metrics import adjusted_rand_score, normalized_mutual_info_score, silhouette_score
from sklearn.cluster import KMeans
import numpy as np

# np.random.seed(0)

D = load_iris()
X_train, y_train = D.data, D.target

centroids = []
for i in range(3):  # 给出三个不同类别的中心点作为簇的初始中心点
    for j, y in enumerate(y_train):
        if i == y:
            centroids.append(X_train[j])
            break
centroids = np.array(centroids)

dict_title2init_option = {'普通K-Means': 'random', '指定了初始聚类中心的K-Means': centroids, 'K-Means++': 'k-means++'}

for title, init_option in dict_title2init_option.items():
    print(title, 'with datasets iris')
    kmeans = KMeans(n_clusters=3, init=init_option)
    kmeans.fit(X_train)
    labels = kmeans.labels_
    ARI = adjusted_rand_score(y_train, labels)
    print('\t ARI = ', ARI)
    NMI = normalized_mutual_info_score(y_train, labels)
    print('\t NMI = ', NMI)

    print('\t 轮廓系数：', silhouette_score(X_train, labels))
